# 导入库
import json
from tqdm import tqdm
from ..config.config import Config
from py2neo import Graph, Node, Relationship
from py2neo import NodeMatcher, RelationshipMatcher
import pdb
# 导入模型预测方法
from ..model_train_predict.lstm_model_train import LSTMModelTrainPredict

class NeoGraph():
    def __init__(self):
        self.config = Config()
        self.url = "http://localhost:7474/"
        self.username = "neo4j"
        self.password = "324820"
        self.lstm_model = LSTMModelTrainPredict()
        self.graph = Graph(self.url, auth=(self.username, self.password))
        self.node_match = NodeMatcher(self.graph)
        self.relation_match = RelationshipMatcher(self.graph)

    def connect_neo4j(self):
        # 实例化图数据库
        # 连接neo4j数据库，输入地址、用户名、密码
        self.graph = Graph(self.url, auth=(self.username, self.password))
        # 实例化节点查询对象
        self.node_match = NodeMatcher(self.graph)
        # 实例化关系查询对象
        self.relation_match = RelationshipMatcher(self.graph)

    # 创建节点
    def create_node(self, label, attrs):
        # 第一次：如果你的图数据库之前的数据是没有用的可以直接清空
        # graph.delete_all()
        # 构建节点的属性，如果一个节点有多个属性要同时判断，所以加上and
        value = self.match_node(label, attrs)
        # breakpoint()
        # pdb.set_trace()  # 在这里设置断点
        # 如果要创建的节点不存在则再创建
        if value is None:
            node = Node(label, attrs)
            node = self.graph.create(node)
            return node
        return value

    # 使用NodeMatcher查询节点
    def match_node(self, label, attrs):
        condition = ""
        for key, value in attrs.items():
            condition += '_.%s=' % key + '\"' + value + '\"' + " and "
        # 将condition最后的 and 字符去掉(注意and前后有两个空格)
        condition = condition[:-5]
        # 根据属性条件查询节点是否已经存在，若存在则返回该节点，否则返回None。
        # 注意加上.first()返回一个节点结果，不加.first()返回所有符合要求的节点结果
        value = self.node_match.match(label).where(condition).first()
        return value

    # 建立两个节点之间的关系
    def create_relationship(self, label1, attrs1, label2, attrs2, r_name):
        value1 = self.match_node(label1, attrs1)
        value2 = self.match_node(label2, attrs2)
        # 判断实体是否均存在，否则无法创建关系
        if value1 is None or value2 is None:
            return False
        # 判断是否已经创建完关系，如果已经创建就不用再重复定义了。当然这一步也可以省略，因为相同实体对创建相同关系时，结果会覆盖
        # rel_value = match_relation(node1=value1, node2=value2, r_type=r_name)
        # print("rel_value:{}".format(rel_value))
        # if rel_value:
        #     return False
        r = Relationship(value1, r_name, value2)
        self.graph.create(r)

    # 获取实体类型以及关系类型
    def get_spo(self):
        ''''
        函数功能描述：
        1、在正常使用模型预测spo时，因为模型不能第一时间给出subject和object的类型，
        因此，这里我们借助训练集中的类型即可
        2、如果训练集也没有，可以根据预测关系的类别自己定义比如：出生日期，subject一般为人名，object一般为日期或时间
        '''
        dict1 = {}
        with open(self.config.bert_train_path, 'r', encoding='utf-8') as fr:
            for line in fr.readlines():
                line = json.loads(line)
                spo_list = line['spo_list']
                for spo in spo_list:
                    if spo['predicate'] not in dict1:
                        dict1[spo['predicate']] = spo
                if len(dict1) == 18:
                    break
        # with open('rel_type.json', 'a', encoding='utf-8') as fw:
        #     fw.write(json.dumps(dict1, ensure_ascii=False, indent=4))
        return dict1

    # 根据模型预测spo,准备数据
    def ready_data(self):
        '''
        通过加载训练好的模型，对需要待预测关系的文本实现预测，并将预测结果存储到文件中
        :return:
        '''
        # 获取包含实体和关系类型的字典
        type_dict = self.get_spo()
        # 导入训练好的模型
        model_path = f'{self.config.save_model_path}\last_model.pth'
        my_model = self.lstm_model.load_model(model_path)
        # 读取测试集文档进行预测
        with open(self.config.bert_test_path, 'r', encoding='utf-8')as fr:
            lines = fr.readlines()
        # 遍历所有测试集
        for line in tqdm(lines):
            line = json.loads(line)
            sample = line['text']
            outputs = self.lstm_model.model2test(sample, my_model)
            if len(outputs) == 0:
                continue
            spo_list = outputs['spo_list']
            # print(spo_list)
            for spo in spo_list:
                if spo['predicate'] in type_dict:
                    spo['subject_type'] = type_dict[spo['predicate']]['subject_type']
                    spo['object_type'] = type_dict[spo['predicate']]['object_type']
                    with open(f'{self.config.entertainment_data_path}\predict_spo.json', 'a', encoding='utf-8') as fw:
                        fw.write(json.dumps(spo, ensure_ascii=False) + '\n')

    # 将SPO数据全部导入Neo4j数据库中
    def load_file_create_map(self):
        # 首先需要运行ready_data函数获取模型预测的spo三元组数据，这里是通过文件形式存储数据（方便反复使用）
        # 第一次运行可以清空所有图数据
        self.graph.delete_all()
        # 第一步获取数据
        with open(f'{self.config.entertainment_data_path}\predict_spo.json', 'r', encoding='utf-8') as fr:
            for line in tqdm(fr.readlines()):
                line = json.loads(line)
                # todo: 定义主实体节点
                # 定义主实体节点的标签
                subject_label = line["subject_type"]
                # 定义主实体节点的属性
                sub_attrs = {'name': line["subject"]}
                self.create_node(subject_label, sub_attrs)

                # todo: 定义客实体节点
                # 定义客实体节点的标签
                object_label = line["object_type"]
                # 定义主实体节点的属性
                if "日期" in line['predicate']:
                    obj_attrs = {'date': line["object"]}
                else:
                    obj_attrs = {'name': line["object"]}
                self.create_node(object_label, obj_attrs)

                # todo: 定义主实体和客实体之间的关系
                r_name = line["predicate"]
                self.create_relationship(subject_label, sub_attrs, object_label, obj_attrs, r_name)

    # 查询Neo4j图数据库数据
    def use_neo4j2search(self):
        # 查询图中节点和关系类型有哪些
        # 1、查询一共有多少种节点类型
        node_labels = self.graph.schema.node_labels
        print(F"Neo4j图数据库中存在的节点标签为：{node_labels}")

        # 2、查询一共有多少种关系类型
        relation_types = self.graph.schema.relationship_types
        print(F"Neo4j图数据库中存在的关系类型为：{relation_types}")

        # 3、查询节点，这里以label="人物", attrs = {"name": "李晨"}为例
        node1 = self.match_node(label='人物', attrs={"name": '李晨'})
        print(node1)

        # 4、查询关系，以上述的名字为“李晨”的节点为出发点，进行关系的查询
        # eg1：查询“李晨”节点的所有关系：先查节点，再查关系，r_type=None表示任意的关系类型都可以
        ship_list1 = list(self.relation_match.match([node1], r_type=None))
        for reship in ship_list1:
            print(reship)
        print('*' * 80)

        # eg2：查询“李晨”和“中国”的关系：两个节点的顺序表示要匹配的关系方向，r_type=None表示任意的关系类型都可以
        node2 = self.match_node(label='国家', attrs={"name": '中国'})
        ship_list12 = list(self.relation_match.match([node1, node2], r_type=None))
        for reship in ship_list12:
            print(reship)
        print('*' * 80)

        # eg3：查询某一类关系：第一个参数为None，第二个参数r_type=指定关系类型
        ship_list3 = list(self.relation_match.match(None, r_type="出生日期"))
        for reship in ship_list3:
            print(reship)
            break